import codecs
import csv

import cv2
import numpy as np

from Utils.DataLoader import IITDLoader

if __name__ == '__main__':
    dl = IITDLoader('/home/yjy/dataset/IITD-ROI-NN')
    # dl = IITDLoader('/home/yjy/dataset/IITD-ROI-Zhang')
    csv_path = '/home/yjy/PycharmProjects/HRNet-Facial-Landmark-Detection/output/IITD/palm_alignment_iitd_hrnet_w18/predictions_all.csv'
    # csv_path = '/home/yjy/dataset/iitd_label.csv'
    csv_path_edited = '/home/yjy/dataset/iitd_label_refine.csv'

    with codecs.open(csv_path, 'r', 'utf-8') as csv_file, codecs.open(csv_path_edited, 'w', 'utf-8') as csv_file_out:
        fieldnames = ['file_name', 'x0', 'x2', 'y0', 'y2']
        reader = csv.DictReader(csv_file)
        writer = csv.DictWriter(csv_file_out, fieldnames=fieldnames)
        writer.writeheader()
        for row in reader:
            fname = row[fieldnames[0]]
            image = dl.get(fname)
            th, image_thr = cv2.threshold(image, 70, 1, cv2.THRESH_BINARY)
            white_num = np.sum(image_thr, keepdims=False)
            print(white_num)
            if white_num > 16380:
                writer.writerow(row)
            else:
                print(fname)
